from sortedcontainers import SortedList
import numpy as np
import json
import random
import math
from pathlib import Path
from datetime import datetime
from concurrent.futures import ThreadPoolExecutor, as_completed
from multiprocessing import Pool, cpu_count, Manager

def _generate_idch_parallel(args):
    """并行生成IDCH染色体的辅助函数"""
    local_search = LocalSearch(
        args['valid_points'],
        args['valid_matrix'],
        args['node_mapping'],
        args['profits'],
        args['time_limit'],
        args['valid_time'],
        args['penalty_rate'],
        args['penalty_interval']
    )
    return local_search.IDCH()

def _process_chromosome_parallel(args):
    """并行处理染色体的辅助函数"""
    chromo, params = args
    local_search = LocalSearch(
        params['valid_points'],
        params['valid_matrix'],
        params['node_mapping'],
        params['profits'],
        params['time_limit'],
        params['valid_time'],
        params['penalty_rate'],
        params['penalty_interval']
    )
    return local_search.LS(chromo)

# 并行处理控制开关
USE_PARALLEL = True  # 设置为False可禁用并行处理
from UAS.algorithm.evolutionary import Selection, Crossover
from UAS.algorithm.local_search import LocalSearch
from UAS.algorithm.chromosome import Chromosome
from UAS.data.excel_parser import ExcelParser

class MemeticAlgorithm:
    def __init__(self, valid_points, valid_matrix, node_mapping, profits, 
                 max_mutation_prob=0.1, min_mutation_prob=0.001,
                 population_size=50, max_iterations=50, callback=None,
                 penalty_interval=10, penalty_rate=10,
                 time_limit=120, valid_time=150,
                 performance_mode=False):
        self.callback = callback
        self.should_terminate = False
        self.performance_mode = performance_mode
        # 获取逻辑核心数
        print(f"USE_PARALLEL={USE_PARALLEL}")
        try:
            self.physical_cores = cpu_count()
            self.physical_cores = 4
            print(f"可用CPU核心数: {self.physical_cores}")
        except:
            self.physical_cores = 1
        """
        初始化memetic算法
        :param valid_points: 有效点位列表
        :param valid_matrix: 有效邻接矩阵
        :param node_mapping: 节点编码映射
        :param profits: 点位收益列表
        :param max_mutation_prob: 最大变异概率
        :param min_mutation_prob: 最小变异概率
        :param population_size: 种群大小
        :param max_iterations: 最大迭代次数
        """
        self.valid_points = valid_points
        self.valid_matrix = valid_matrix
        self.node_mapping = node_mapping
        self.profits = profits
        self.reverse_mapping = {v: k for k, v in node_mapping.items()}
        # 时间参数
        self.time_limit = time_limit  # 时间限制
        self.valid_time = valid_time  # 有效时间
        # 超时惩罚规则
        self.penalty_interval = penalty_interval  # 惩罚间隔
        self.penalty_rate = penalty_rate  # 惩罚值
        # 变异概率参数
        self.max_mutation_prob = max_mutation_prob
        self.min_mutation_prob = min_mutation_prob
        # 种群参数
        self.population_size = population_size
        self.max_iterations = max_iterations

    def generate_random_chromosome(self):
        """生成随机染色体"""
        import random
        n = len(self.valid_points)
        
        # 生成路径部分：0开头，随机选择1的位置
        path_length = random.randint(2, n-1)  # 至少包含0和1
        path_nodes = [0] + random.sample(range(1, n), path_length-1)
        
        # 随机选择1的位置（不能是第一个）
        one_idx = random.randint(1, len(path_nodes)-1)
        path_nodes[one_idx] = 1
        
        # 生成未访问节点部分
        all_nodes = set(range(n))
        visited_nodes = set(path_nodes)
        unvisited_nodes = list(all_nodes - visited_nodes)
        random.shuffle(unvisited_nodes)
        
        # 组合完整序列
        sequence = path_nodes + unvisited_nodes
        
        return Chromosome(sequence, self.valid_points, self.profits, self.valid_matrix,
                         self.time_limit, self.valid_time,
                         self.penalty_interval, self.penalty_rate)
        
    def save_config(self, output_dir='output'):
        """保存算法配置"""
        Path(output_dir).mkdir(exist_ok=True)
        
        # 保存valid_points
        with open(f'{output_dir}/valid_points.json', 'w') as f:
            json.dump(self.valid_points, f)
            
        # 保存valid_matrix
        np.save(f'{output_dir}/valid_matrix.npy', self.valid_matrix)
        
        # 保存node_mapping
        with open(f'{output_dir}/node_mapping.json', 'w') as f:
            json.dump(self.node_mapping, f)
            
        # 保存profits
        with open(f'{output_dir}/profits.json', 'w') as f:
            json.dump(self.profits, f)
            
    @classmethod
    def load_config(cls, output_dir='output'):
        """加载算法配置"""
        # 加载valid_points
        with open(f'{output_dir}/valid_points.json') as f:
            valid_points = json.load(f)
            
        # 加载valid_matrix
        valid_matrix = np.load(f'{output_dir}/valid_matrix.npy')
        
        # 加载node_mapping
        with open(f'{output_dir}/node_mapping.json') as f:
            node_mapping = json.load(f)
            
        # 加载profits
        with open(f'{output_dir}/profits.json') as f:
            profits = json.load(f)
            
        return cls(valid_points, valid_matrix, node_mapping, profits)
    
    def _sort_key(self, chromosome):
        """SortedList排序键函数"""
        return (-chromosome.total_profit, chromosome.total_cost)

    def evaluate_population(self, population):
        """评估种群质量"""
        best = population[0]
        worst = population[-1]
        avg_profit = sum(c.total_profit for c in population) / len(population)
        avg_cost = sum(c.total_cost for c in population) / len(population)
        
        # 获取前3个染色体
        top3 = population[:3]
        top3_chromosomes = [
            {
                'sequence': c.sequence,
                'profit': c.total_profit,
                'cost': c.total_cost
            }
            for c in top3
        ]
        
        return {
            'best_profit': best.total_profit,
            'best_cost': best.total_cost,
            'worst_profit': worst.total_profit,
            'worst_cost': worst.total_cost,
            'avg_profit': avg_profit,
            'avg_cost': avg_cost,
            'valid_count': sum(1 for c in population if c.total_cost <= self.valid_time),
            'top3_chromosomes': top3_chromosomes
        }

    def record_statistics(self, population, iteration=0):
        """记录种群统计信息"""
        stats = self.evaluate_population(population)
        stats['iteration'] = iteration
        stats['timestamp'] = datetime.now().isoformat()
        
        # 保存到文件
        output_dir = Path('output/stats')
        output_dir.mkdir(exist_ok=True)
        with open(output_dir / f'stats_{iteration}.json', 'w') as f:
            json.dump(stats, f)
            
        # 打印前3个染色体信息
        print(f"\nTop 3 Chromosomes in iteration {iteration}:")
        for i, chromo in enumerate(stats['top3_chromosomes'], 1):
            #print(f"  {i}. Sequence: {chromo['sequence']}")
            print(f"{i}     Profit: {chromo['profit']}, Cost: {chromo['cost']}")
        
        return stats

    def visualize_population(self, population, iteration=0):
        """可视化种群分布
        :param population: 当前种群
        :param iteration: 当前迭代次数
        """
        import matplotlib.pyplot as plt
        
        # 准备数据
        profits = [c.total_profit for c in population]
        costs = [c.total_cost for c in population]
        
        # 创建图表
        plt.figure(figsize=(12, 6))
        
        # 散点图
        plt.subplot(1, 2, 1)
        plt.scatter(costs, profits, alpha=0.6)
        plt.xlabel('Cost')
        plt.ylabel('Profit')
        plt.title(f'Population Distribution (Iteration {iteration})')
        
        # 箱线图
        plt.subplot(1, 2, 2)
        plt.boxplot([profits, costs], labels=['Profit', 'Cost'])
        plt.title(f'Population Statistics (Iteration {iteration})')
        
        # 保存图像
        output_dir = Path('output/plots')
        output_dir.mkdir(exist_ok=True)
        plt.savefig(output_dir / f'population_{iteration}.png')
        plt.close()

    def _process_idch(self, local_search):
        """处理IDCH染色体生成"""
        return local_search.IDCH()

    def initialize_population(self):
        """初始化种群"""
        population = SortedList(key=self._sort_key)
        
        # 计算IDCH染色体数量（至少4个）
        idch_count = max(4, int(self.population_size * 0.05))

        # 并行生成IDCH染色体
        if USE_PARALLEL:
            # 使用multiprocessing.Pool并行执行IDCH
            with Pool(processes=self.physical_cores) as pool:
                # 创建LocalSearch实例的参数
                args = {
                    'valid_points': self.valid_points,
                    'valid_matrix': self.valid_matrix,
                    'node_mapping': self.node_mapping,
                    'profits': self.profits,
                    'time_limit': self.time_limit,
                    'valid_time': self.valid_time,
                    'penalty_rate': self.penalty_rate,
                    'penalty_interval': self.penalty_interval
                }
                
                # 并行生成IDCH染色体
                results = pool.map(_generate_idch_parallel, [args]*idch_count)
                for chromo in results:
                    population.add(chromo)
        else:
            # 串行生成IDCH染色体
            local_search = LocalSearch(self.valid_points, self.valid_matrix, 
                                    self.node_mapping, self.profits,
                                    self.time_limit, self.valid_time,
                                    self.penalty_rate, self.penalty_interval)
            for _ in range(idch_count):
                chromo = self._process_idch(local_search)
                population.add(chromo)

        # 生成剩余随机染色体
        for _ in range(self.population_size - idch_count):
            chromosome = self.generate_random_chromosome()
            population.add(chromosome)
        
        return population

    @staticmethod
    def _parallel_idch_wrapper(args):
        """IDCH生成的包装函数(用于并行化)"""
        local_search = args
        return local_search.IDCH()

    def format_population_results(self, population, top_n=3, iteration=None):
        """
        格式化种群结果，将编码转换为实际点位编号，并计算实际时间
        :param population: 种群列表
        :param top_n: 要输出的前N个染色体
        :param iteration: 当前迭代次数
        :return: 包含格式化结果的列表
        """
        results = []
        # 根据模式调整top_n
        if self.performance_mode and iteration is not None and iteration != self.max_iterations - 1:
            return []  # 高性能模式下非最后迭代不保存结果
            
        for i, chromosome in enumerate(population[:top_n]):
            # 找到第一个1的位置
            try:
                end_idx = chromosome.sequence.index(1)
            except ValueError:
                end_idx = len(chromosome.sequence) - 1
                
            # 截取从0到1的路径
            truncated_sequence = chromosome.sequence[:end_idx+1]
            
            # 将路径编码转换为实际点位编号
            path = [self.reverse_mapping[code] for code in truncated_sequence]

            #print(f"path={path}")
            
            # 计算各点收益
            point_profits = [self.profits[code] for code in truncated_sequence]
            
            # 计算路径时间（将矩阵值除以10）
            path_times = []
            total_time = 0
            for j in range(len(truncated_sequence)-1):
                from_node = truncated_sequence[j]
                to_node = truncated_sequence[j+1]
                time = self.valid_matrix[from_node][to_node]
                path_times.append(round(time/10, 1)) #路径时间是实际的时间
                total_time += time
            total_time = int(total_time) #总时间是整数，还是乘以10的状态
            
            # 计算惩罚后的总收益
            if total_time <= self.valid_time:
                over_time = max(0, total_time - self.time_limit)
                penalty = math.ceil(over_time / self.penalty_interval) * self.penalty_rate
                total_profit = sum(point_profits) - penalty
                #print(f"total_time={total_time}, over_time={over_time}, penalty={penalty}")
            else:
                total_profit = 0
            total_time = round(total_time/10, 1) #将总时间转换为实际时间

            # 格式化结果
            result = {
                'rank': i+1,
                'path': path,
                'point_profits': point_profits,
                'total_profit': total_profit,  # 应用惩罚后的总收益
                'path_times': path_times,
                'total_time': total_time  # 实际总时间
            }
            results.append(result)

        # 构建文件名
        if iteration is not None:
            filename = f'results/results_{iteration}.json'
        else:
            filename = 'results/results_final.json'
        
        # 保存结果到 JSON 文件
        results_data = {
            'iteration': iteration,
            'timestamp': datetime.now().isoformat(),
            'results': results
        }
        with open(filename, 'w') as f:
            json.dump(results_data, f, indent=4)
        
        return results

    def _process_chromosome(self, chromo):
        """处理单个染色体的局部搜索(用于并行化)"""
        local_search = LocalSearch(self.valid_points, self.valid_matrix,
                                 self.node_mapping, self.profits,
                                 self.time_limit, self.valid_time,
                                 self.penalty_rate, self.penalty_interval)
        return local_search.LS(chromo)

    def terminate(self):
        """终止算法"""
        self.should_terminate = True

    def run(self):
        """主循环"""
        if self.callback:
            self.callback({"status": "start", "message": "Algorithm started"})
        
        # 重置终止标志
        self.should_terminate = False
        # 初始化种群
        population = self.initialize_population()
        # 初始化局部搜索相关变量
        iteration_without_improvement = 0
        previous_best = population[0].total_profit if population else 0
        
        # 主循环
        for iteration in range(self.max_iterations):
            # 初始化选择器并重置选择状态
            selection = Selection(population)
            selection.reset_selection()
            
            new_population = SortedList(key=self._sort_key)

            print(f"Start LOX crossover for iteration {iteration}")

            while True:
                # 选择父代
                if len(population) < int(self.population_size * 0.75):
                    p1, p2 = selection.random_select()
                else:
                    p1, p2 = selection.tournament_select()
                    
                if p1 is None:  # 可用个体不足
                    break

                # 交叉生成子代
                crossover = Crossover(self.valid_points)
                c1, c2 = crossover.LOX(p1.sequence, p2.sequence)
                
                # 评估子代
                child1 = Chromosome(c1, self.valid_points, self.profits, self.valid_matrix,
                                  self.time_limit, self.valid_time,
                                  self.penalty_interval, self.penalty_rate)
                child2 = Chromosome(c2, self.valid_points, self.profits, self.valid_matrix,
                                  self.time_limit, self.valid_time,
                                  self.penalty_interval, self.penalty_rate)
                
                # 检查子代是否相似
                def _is_similar(c1, c2, threshold=0.8):
                    """检查两个染色体是否相似"""
                    # 检查profit和cost
                    if abs(c1.total_profit - c2.total_profit) > 0:
                        return False
                    if abs(c1.total_cost - c2.total_cost) > 2:
                        return False
                    
                    # 计算序列相似度
                    match_count = sum(1 for a,b in zip(c1.sequence, c2.sequence) if a == b)
                    similarity = match_count / len(c1.sequence)
                    
                    return similarity >= threshold

                # 计算当前相似度阈值
                if iteration < 0.9 * self.max_iterations:
                    current_threshold = 0.8 + (0.2 * (iteration / (0.9 * self.max_iterations)))
                else:
                    current_threshold = 1

                # 添加子代到新种群
                skip_child1 = any(_is_similar(child1, c, current_threshold) for c in new_population)
                if not skip_child1:
                    new_population.add(child1)
                    
                skip_child2 = any(_is_similar(child2, c, current_threshold) for c in new_population)
                if not skip_child2:
                    new_population.add(child2)
                    
                # 添加父代到新种群
                skip_p1 = any(_is_similar(p1, c, current_threshold) for c in new_population)
                if not skip_p1:
                    new_population.add(p1)
                    
                skip_p2 = any(_is_similar(p2, c, current_threshold) for c in new_population)
                if not skip_p2:
                    new_population.add(p2)
            
            # 确保新种群不为空
            if len(new_population) == 0:
                # 如果新种群为空，重新初始化
                print("Warning: Population became empty, reinitializing")
                population = self.initialize_population()
            else:
                # 用新种群替换旧种群
                population = new_population

            # 初始化局部搜索
            local_search = LocalSearch(self.valid_points, self.valid_matrix,
                                     self.node_mapping, self.profits,
                                     self.time_limit, self.valid_time,
                                     self.penalty_rate, self.penalty_interval)
            
            # 计算变异概率
            if iteration < (2/3)*self.max_iterations:
                # 递减阶段
                pm = self.max_mutation_prob - (self.max_mutation_prob - self.min_mutation_prob) * \
                     (iteration / ((2/3)*self.max_iterations))
            else:
                # 递增阶段
                pm = self.min_mutation_prob + ((self.max_mutation_prob + self.min_mutation_prob)/2 - self.min_mutation_prob) * \
                     ((iteration - (2/3)*self.max_iterations) / ((1/3)*self.max_iterations))

            print(f"pm={pm:.3f}")
            
            # 局部搜索优化(根据USE_PARALLEL选择串行或并行)
            chromo_to_optimize = [chromo for chromo in population 
                                if random.random() < pm]
            
            if USE_PARALLEL:
                # 使用multiprocessing.Pool并行处理
                with Pool(processes=self.physical_cores) as pool:
                    # 准备参数
                    args_list = [(chromo, {
                        'valid_points': self.valid_points,
                        'valid_matrix': self.valid_matrix,
                        'node_mapping': self.node_mapping,
                        'profits': self.profits,
                        'time_limit': self.time_limit,
                        'valid_time': self.valid_time,
                        'penalty_rate': self.penalty_rate,
                        'penalty_interval': self.penalty_interval
                    }) for chromo in chromo_to_optimize]
                    
                    # 并行处理
                    results = pool.map(_process_chromosome_parallel, args_list)
                    
                    # 更新种群
                    for original, optimized in zip(chromo_to_optimize, results):
                        population.pop(population.index(original))
                        population.add(optimized)
            else:
                for chromo in chromo_to_optimize:
                    optimized = self._process_chromosome(chromo)
                    population.pop(population.index(chromo))
                    population.add(optimized)

            # 记录统计信息
            if len(population) > 0:
                if not self.performance_mode:
                    stats = self.record_statistics(population, iteration)
                    self.visualize_population(population, iteration)
                else:
                    # 高性能模式只记录基本信息
                    if iteration % 20 == 0 or iteration == self.max_iterations - 1:
                        stats = self.evaluate_population(population)
                        print(f"Iteration {iteration}: Best Profit={stats['best_profit']}")
            else:
                print("Error: Population is empty")
                continue
            
            # 更新无改进迭代计数
            current_best = stats['best_profit']
            if current_best <= previous_best:
                iteration_without_improvement += 1
                # 如果连续n/2次没有改进，进行多样化处理
                if iteration_without_improvement >= len(self.valid_points) // 2:
                    # 计算要灭绝的染色体数量（种群大小的1/3取整）
                    num_to_extinct = len(population) // 3
                    
                    # 选择total_profit最小的染色体
                    to_extinct = population[-num_to_extinct:]  # 取最后num_to_extinct个（按收益排序）
                    
                    # 打印日志
                    print(f"Extincting {num_to_extinct} chromosomes and diversifying")
                    
                    # 移除选中的染色体
                    for chromo in to_extinct:
                        population.remove(chromo)
                    
                    # 生成1倍于灭绝数量的新染色体
                    num_new = 1 * num_to_extinct
                    new_chromosomes = [self.generate_random_chromosome() for _ in range(num_new)]
                    
                    # 对其中3个新染色体进行局部搜索优化
                    to_optimize = random.sample(new_chromosomes, min(3, num_new))
                    optimized_chromosomes = [local_search.LS(chromo) for chromo in to_optimize]
                    
                    # 将所有新染色体加入种群
                    for chromo in new_chromosomes:
                        population.add(chromo)
                    for chromo in optimized_chromosomes:
                        population.add(chromo)
                    
                    # 重置计数器
                    iteration_without_improvement = 0
            else:
                iteration_without_improvement = 0
            previous_best = current_best
            
            # 打印进度并触发回调
            #progress_msg = f"Iteration {iteration}: Best Profit={stats['best_profit']}, Best Cost={stats['best_cost']}, Valid Count={stats['valid_count']}"
            #print(progress_msg)
            
            # 每1次迭代保存一次结果
            #self.format_population_results(population, top_n=1, iteration=iteration)
                
            # if self.callback:
            #     self.callback({
            #         "iteration": iteration,
            #         "best_profit": stats['best_profit'],
            #         "best_cost": stats['best_cost'],
            #         "valid_count": stats['valid_count'],
            #         "message": progress_msg
            #     })
            
            # 检查终止标志
            if self.should_terminate:
                print("算法被用户终止")
                break

            # 终止条件：连续n^2代没有改进
            n = len(self.valid_points)
            if iteration > n**2 and stats['best_profit'] == self.evaluate_population(population[:n**2])['best_profit']:
                break
                
        # 保存最终结果
        self.format_population_results(population, top_n=5, iteration=None)
        return population

if __name__ == "__main__":
    import sys
    from pathlib import Path
    # 添加项目根目录到Python路径
    sys.path.append(str(Path(__file__).parent.parent.parent))
    from UAS.data.graph_builder import GraphBuilder
    from UAS.data.excel_parser import ExcelParser
    
    # 初始化解析器和图构建器
    parser = ExcelParser("d:/Work/Projects/UAS/output/valid_points.xlsx", 
                        "d:/Work/Projects/UAS/output/valid_matrix.xlsx")
    parser.load_data()
    builder = GraphBuilder(parser)
    
    # 获取graph_builder的数据
    valid_points, profits, valid_matrix, node_mapping = builder.build_graph()
    
    # 打印验证数据
    print("从graph_builder获取的数据：")
    print(f"有效点位数量: {len(valid_points)}")
    print(f"有效点位列表: {valid_points}")
    print(f"\n点位收益: {profits}")
    print(f"\n邻接矩阵形状: {valid_matrix.shape}")
    print("邻接矩阵内容:")
    print(valid_matrix)
    print(f"\n节点编码映射数量: {len(node_mapping)}")
    print("节点编码映射:")
    for point, code in node_mapping.items():
        print(f"  {point}: {code}")
        
    # 初始化memetic算法
    ma = MemeticAlgorithm(valid_points, valid_matrix, node_mapping, profits)
    
    # 生成并打印随机染色体
    print("\n随机生成的染色体:")
    random_chromosome = ma.generate_random_chromosome()
    print(f"序列: {random_chromosome.sequence}")
    print(f"总收益: {random_chromosome.total_profit}")
    print(f"总成本: {random_chromosome.total_cost}")
    print(f"是否有效: {'是' if random_chromosome.total_cost <= ma.valid_time else '否'}")

    # 运行主循环
    final_population = ma.run()
    
    # 输出最终结果
    print("\n最终种群统计:")
    print(f"种群大小: {len(final_population)}")
    
    # 输出格式化后的种群结果
    print("\n=== Top 3 Solutions ===")
    formatted_results = ma.format_population_results(final_population, top_n=3)
    for result in formatted_results:
        print(f"\nRank {result['rank']}:")
        print("  Path:")
        print("    " + " -> ".join(result['path']))
        print("  Point Profits:")
        print("    " + ", ".join(f"{p}" for p in result['point_profits']))
        print("  Path Times:")
        print("    " + ", ".join(f"{t:.1f} min" for t in result['path_times']))
        print(f"  Total Time: {result['total_time']:.1f} min")
        print(f"  Total Profit: {result['total_profit']}")
    
    # 输出最佳解决方案摘要
    best = formatted_results[0]
    print("\n=== Best Solution Summary ===")
    print(f"Total Profit: {best['total_profit']}")
    print(f"Total Time: {best['total_time']:.1f} min")
    print("Path:")
    print("  " + " -> ".join(best['path']))
    
    print(f"\n最佳染色体: {final_population[0]}")
    print(f"最佳染色体收益: {final_population[0].total_profit}")
    print(f"最佳染色体成本: {final_population[0].total_cost}")
